#include "execution/executors/topn_executor.h"

namespace bustub {

TopNExecutor::TopNExecutor(ExecutorContext *exec_ctx, const TopNPlanNode *plan,
                           std::unique_ptr<AbstractExecutor> &&child_executor)
    : AbstractExecutor(exec_ctx), plan_(plan), child_executor_(std::move(child_executor)) {}

void TopNExecutor::Init() {
  child_executor_->Init();
  TopNode t;
  RID rid;
  while (child_executor_->Next(&t.tuple_, &rid)) {
    for (const auto &order_by : plan_->order_bys_) {
      t.vec_.emplace_back(std::pair{order_by.first, order_by.second->Evaluate(&t.tuple_, plan_->OutputSchema())});
    }
    priority_queue_.push(t);
    if (priority_queue_.size() > plan_->GetN()) {
      priority_queue_.pop();
    }
    t = TopNode{};
  }
  while (!priority_queue_.empty()) {
    auto entity = priority_queue_.top();
    priority_queue_.pop();
    final_res_.emplace_back(entity.tuple_);
  }
}

auto TopNExecutor::Next(Tuple *tuple, RID *rid) -> bool {
  if (!final_res_.empty()) {
    auto entity = final_res_.back();
    final_res_.pop_back();
    *tuple = entity;
    *rid = tuple->GetRid();
    return true;
  }
  return false;
}

auto TopNExecutor::GetNumInHeap() -> size_t { return priority_queue_.size(); };

}  // namespace bustub
